import numpy as np
import copy


class OUActionNoise:  # OUActionNoise类用于生成 Ornstein-Uhlenbeck 过程的噪声
    def __init__(self, size, mu=0., theta=0.15, sigma=0.2):
        self.size = size
        self.mu = mu * np.ones(size)
        self.theta = theta
        self.sigma = sigma
        self.reset()

    def __call__(self, episode=None, max_episodes=None, size=None):
        """
        Generate Ornstein-Uhlenbeck noise with optional episode-based decay.

        :param episode: 当前的 episode 数，用于噪声衰减
        :param max_episodes: 最大的训练 episode 数，用于噪声衰减
        :param size: 如果传入，动态改变噪声的维度
        :return: 生成的噪声
        """
        # 如果没有传递 size 参数，默认使用初始化时的 size
        if size is None:
            size = self.size
        elif size != len(self.state):
            # 如果 size 发生变化，调整 mu 和 state
            self.mu = np.zeros(size)  # 调整 mu 为新 size 的大小
            self.reset()

        # 计算基于 episode 的噪声衰减（线性衰减）
        if episode is not None and max_episodes is not None:
            # 使用线性衰减
            decayed_sigma = self.sigma * (1 - episode / max_episodes)
        else:
            decayed_sigma = self.sigma  # 如果没有衰减信息，则使用初始的标准差

        # 使用 Ornstein-Uhlenbeck 过程生成噪声
        dx = self.theta * (self.mu - self.state) + decayed_sigma * np.random.randn(size)
        self.state = self.state + dx
        return self.state

    def reset(self):
        """Reset the internal state (= noise) to mean (mu)."""
        self.state = copy.copy(self.mu)


class NormalNoise:  # NormalNoise 类用于生成服从正态分布的噪声
    def __init__(self, normal_scalar):
        self.normal_scalar = normal_scalar

    def __call__(self):
        return np.random.randn(1) * self.normal_scalar




